import random
from .data_processor import DataProcessor
from .feature_extractor import FeatureExtractor
from .template_manager import TemplateManager

class QAGenerator:
    def __init__(self, cycle_data, step_data, record_data, cycle_group_indices, charge_discharge_rates, charge_discharge_currents, reactant_mass):
        self.cycle_data = cycle_data
        self.step_data = step_data
        self.record_data = record_data
        self.cycle_group_indices = cycle_group_indices
        self.charge_discharge_rates = charge_discharge_rates
        self.charge_discharge_currents = charge_discharge_currents
        self.reactant_mass = reactant_mass

    def generate_capacity_attenuation_trend(self, discharge_group_cap, cycle):
        #cycle_ten_list = []
        if cycle // 10 > 3 and cycle // 10 <= len(discharge_group_cap) - 4:
            #print(f"cycle: {cycle}QQQQQQ")
            if (cycle // 10) * 10 == cycle:
                discharge_caps = f"第{(cycle // 10 - 3) * 10}轮、{(cycle // 10 - 2) * 10}轮、{(cycle // 10 - 1) * 10}轮、{(cycle // 10 + 1) * 10}轮、{(cycle // 10 + 2) * 10}轮、{(cycle // 10 + 3) * 10}轮"\
                f"循环的放电容量依次为{discharge_group_cap[cycle // 10 - 3]}Ah、{discharge_group_cap[cycle // 10 - 2]}Ah、{discharge_group_cap[cycle // 10 - 1]}Ah、"\
                f"{discharge_group_cap[cycle // 10 + 1]}Ah、{discharge_group_cap[cycle // 10 + 2]}Ah、{discharge_group_cap[cycle // 10 + 3]}Ah"
            else:
                discharge_caps = f"第{(cycle // 10 - 3) * 10}轮、{(cycle // 10 - 2) * 10}轮、{(cycle // 10 - 1) * 10}轮、{(cycle // 10)* 10}轮、{(cycle // 10 + 1) * 10}轮、{(cycle // 10 + 2) * 10}轮、{(cycle // 10 + 3) * 10}次"\
                f"循环的放电容量依次为{discharge_group_cap[cycle // 10 - 3]}Ah、{discharge_group_cap[cycle // 10 - 2]}Ah、{discharge_group_cap[cycle // 10 - 1]}Ah、{discharge_group_cap[cycle // 10]}Ah、"\
                f"{discharge_group_cap[cycle // 10 + 1]}Ah、{discharge_group_cap[cycle // 10 + 2]}Ah、{discharge_group_cap[cycle // 10 + 3]}Ah"
        elif cycle // 10 <= 3:
            cyc = 1
            index = 6
            discharge_cap_index = "第"
            discharge_cap_values = "循环的放电容量依次为"
            while cyc < cycle // 10:
                #cycle_ten_list.append(cyc * 10)
                #print(f"cyc: {cyc}, cycle: {cycle}")
                #print(f"cycle: {cycle}, cyc: {cyc}")
                discharge_cap_index += f"{cyc * 10}轮、"
                #print(f"discharge_cap_index: {discharge_cap_index}")
                discharge_cap_values += f"{discharge_group_cap[cyc]}Ah、"
                cyc += 1
                index -= 1

            cyc = cycle // 10 + 1
            while index >= 0:
                #cycle_ten_list.append(cyc * 10)
                if (cycle // 10) * 10 != cycle and cyc == cycle // 10 + 1 and cycle > 10:
                    #cycle_ten_list.append((cycle // 10) * 10)
                    discharge_cap_index += f"{(cycle // 10) * 10}轮、"
                    discharge_cap_values += f"{discharge_group_cap[cycle // 10]}(Ah)、"
                discharge_cap_index += f"{cyc * 10}轮"
                #print(f"cycle: {cycle}, cyc: {cyc}")
                discharge_cap_values += f"{discharge_group_cap[cyc]}Ah"
                if index:
                    discharge_cap_index += '、'
                    discharge_cap_values += '、'
                cyc += 1
                index -= 1
            discharge_caps = discharge_cap_index + discharge_cap_values
        else:
            cyc = (cycle // 10) - (6 - (len(discharge_group_cap) - 1 - cycle // 10))
            #print(f"cycle: {cycle}，cyc: {cyc}")
            discharge_cap_index = "第"
            discharge_cap_values = "循环的放电容量依次为"
            while cyc < cycle // 10:
                #print(f"cyc: {cyc}前")
                #cycle_ten_list.append(cyc * 10)
                discharge_cap_index += f"{cyc * 10}次、"
                discharge_cap_values += f"{discharge_group_cap[cyc]}(Ah)、"
                if (cycle // 10) * 10 != cycle and cyc == cycle // 10 - 1:
                    #cycle_ten_list.append((cycle // 10) * 10)
                    discharge_cap_index += f"{(cycle // 10) * 10}次、"
                    discharge_cap_values += f"{discharge_group_cap[cycle // 10]}(Ah)、"
                cyc += 1

            cyc = cycle // 10 + 1
            while cyc < len(discharge_group_cap):
                #cycle_ten_list.append(cyc * 10)
                #print(f"discharge_group_cap[{cyc}]: {discharge_group_cap[cyc]}")
                #print(f"cyc: {cyc}后")
                discharge_cap_index += f"{cyc * 10}次"
                discharge_cap_values += f"{discharge_group_cap[cyc]}Ah"
                if cyc != len(discharge_group_cap) - 1:
                    discharge_cap_index += '、'
                    discharge_cap_values += '、'
                cyc += 1
            discharge_caps = discharge_cap_index + discharge_cap_values
        return discharge_caps

    def generate_dynamic_question(self, cycle):
        """生成动态问题"""
        # offset_list = []
        # printarr = []
        # for i in range(5):
            # offset = random.randint(cycle + 1, cycle + 5)
            # printarr.append(offset)
            # if offset in offset_list or (offset == (cycle // 10) * 10 and cycle // 10 * 10 != cycle) or offset == (cycle // 10 + 1) * 10: #因为随机范围设置中不包括cycle
                # pass
            # else:
                # offset_list.append(offset)
        # offset_list.sort()
        #print(f"cycle: {cycle}, offset_list: {offset_list}")
        #print(f"printarr: {printarr}")
# 
        # question = "问"
        # for j in range(len(offset_list)):
            # question += f"第{offset_list[j]}轮"
            # if j != len(offset_list) - 1:
                # question += '、'
        # question += "的放电量分别是多少？"
        start_cycle_offset = random.randint(1, 5)
        end_cycle_offset = random.randint(1, 5)
        #print(f"\ncycle:{cycle}, start_cycle_offset:{start_cycle_offset}, end_cycle_offset:{end_cycle_offset}\n")

        answer_content = ""
        if start_cycle_offset < end_cycle_offset:
            dynamic_question = f"问第{cycle + start_cycle_offset}到第{cycle + end_cycle_offset}轮的放电量各是多少？"
            for i in range(start_cycle_offset, end_cycle_offset + 1):
                input_answers = {"discharge_cap": self.cycle_data.get("放电容量(Ah)")[cycle + i]}
                answer_content += f"电池在第{cycle + i}轮放电量为$\boxed{{{input_answers['discharge_cap']}}}(Ah)"

                if i != end_cycle_offset:
                    answer_content += '，'
                else:
                    answer_content += '。'
        elif start_cycle_offset > end_cycle_offset:
            dynamic_question = f"问第{cycle + end_cycle_offset}到第{cycle + start_cycle_offset}轮的放电量各是多少？"
            for i in range(end_cycle_offset, start_cycle_offset + 1):
                input_answers = {"discharge_cap": self.cycle_data.get("放电容量(Ah)")[cycle + i]}
                answer_content += f"电池在第{cycle + i}轮放电量为$\boxed{{{input_answers['discharge_cap']}}}(Ah)"
                if i != start_cycle_offset:
                    answer_content += "，"
                else:
                    answer_content += "。"
        else:
            dynamic_question = f"问第{cycle + start_cycle_offset}轮的放电量是多少？"
            input_answers = {"discharge_cap": self.cycle_data.get("放电容量(Ah)")[cycle + start_cycle_offset]}
            answer_content = f"电池在第{cycle + start_cycle_offset}放电容量为$\boxed{{{input_answers["discharge_cap"]}}}(Ah)。"
        #answer = self._generate_range_answer(offset_list)
        return dynamic_question, answer_content

    # def _generate_range_answer(self, offset_list):
        # """生成范围答案"""
        # answer_parts = "电池在"
        # for i in range(len(offset_list)):
            # discharge_cap = self.cycle_data.get("放电容量(Ah)")[offset_list[i] - 1]
            # answer_parts += f"第{offset_list[i]}轮放电量为$\\boxed{{discharge_cap}}$(Ah)"
            # if i != len(offset_list) - 1:
                # answer_parts += '、'
            # else:
                # answer_parts += '。'
        # return answer_parts

    def _generate_single_answer(self, cycle, offset):
        """生成单个答案"""
        discharge_cap = self.cycle_data.get("放电容量(Ah)", [None] * (cycle + offset + 1))[cycle + offset]
        return f"电池在第{cycle + 1 + offset}放电容量为$\\boxed{{discharge_cap}}$(Ah)。"

    def generate_for_cycle(self, cycle_index, cycle):
        """为单个循环生成问答对"""
        time_diff = DataProcessor.get_cycle_time_info(self.step_data, cycle_index)
        if time_diff is None:
            return []
            
        group_index = DataProcessor.get_cycle_group_index(self.cycle_group_indices, cycle)
        if group_index == -1:
            return []
        
        #获取整十的放电容量
        discharge_group_cap = DataProcessor.get_discharge_group_cap(self.cycle_data)
        discharge_caps = self.generate_capacity_attenuation_trend(discharge_group_cap, cycle)

        features = FeatureExtractor.extract_features(self.cycle_data, cycle_index)
        features.update({
            "reactant_mass": self.reactant_mass,
            "cc_charge_ratio": self.charge_discharge_rates[group_index][0],
            "cc_discharge_ratio": self.charge_discharge_rates[group_index][1],
            "cc_charge_current": self.charge_discharge_currents[group_index][0],
            "cc_discharge_current": self.charge_discharge_currents[group_index][1]
        })
        
        #discharge_voltage_group, position = DataProcessor.get_discharge_voltage_group_cap(self.record_data, position)

        dynamic_question, answer = self.generate_dynamic_question(cycle)
        templates = TemplateManager.generate_templates(cycle, features, time_diff, dynamic_question, discharge_caps)
        
        return [{"input": t, "output": answer} for t in templates]

    def generate_all(self):
        """生成所有问答对"""
        qa_pairs = []
        #print(f"{self.cycle_data.shape[0] - 6}&&&&&&&&&&")
        position = 0
        for cycle_index in range(1, self.cycle_data.shape[0] - 6):#索引0对应循环1，索引1对应循环2，所以开始是循环2的数据, 由于for循环设置终止范围时得多加1，所以是.shape[0] - 7
            cycle = cycle_index + 1
            #print(self.cycle_data[self.cycle_data["循环号"] == 1])
            qa_pairs.extend(self.generate_for_cycle(cycle_index, cycle))
            #print(qa_pairs)
        return qa_pairs                                                                                                                                                                                                                                                                                                                       